//============================================================================
// Name        : hmm_test.cpp
// Author      : Long
// Version     :
// Copyright   : Your copyright notice
// Description :
//============================================================================

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <armadillo>

#include "hmm.h"
#include "hmm_partial.h"
#include "dhmm.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

int n, d;
std::string gin, fin, out;

void
process_options(int argc, char** argv);

void
usage(int exitcode = 0)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void
hmm_test()
{
  std::cout << "hmm-test ..." << std::endl;
  hmm model(2, 4);
  arma::mat a(2, 2), b(2, 4);
  arma::vec pi(2);

  pi[0] = 0.5;
  pi[1] = 0.5;

  a.at(0, 0) = 0.7;
  a.at(0, 1) = 0.3;
  a.at(1, 0) = 0.2;
  a.at(1, 1) = 0.8;

  b.at(0, 0) = 0.1;
  b.at(0, 1) = 0.2;
  b.at(0, 2) = 0.4;
  b.at(0, 3) = 0.4;
  b.at(1, 0) = 0.5;
  b.at(1, 1) = 0.3;
  b.at(1, 2) = 0.1;
  b.at(1, 3) = 0.1;

  model.setStart(pi);
  model.setTransition(a);
  model.setEmission(b);
  //model.dump();

  hmm::seq_type sseq, oseq, mapseq;
  model.genSeq(300, sseq, oseq);
  //std::cout << "seq="; hmm::print_seq(sseq, oseq);

  //double prob = model.probability(oseq);
  //std::cout << "logprob=" << prob << " prob=" << exp(prob) << std::endl;

  model.map(oseq, mapseq);
  //std::cout << "mapseq="; hmm::print_seq(mapseq, sseq);
  double error = 0;
  for (unsigned int t = 0; t < sseq.size(); t++)
    if (sseq[t] != mapseq[t]) error++;
  std::cout << "error=" << error << " = " << error / sseq.size() << "%"
      << std::endl;

  std::vector<hmm::seq_type> train_data;
  int N = 1, T = 100;
  for (int i = 0; i < N; i++)
    {
      hmm::seq_type sseq, oseq;
      model.genSeq(T, sseq, oseq);
      train_data.push_back(oseq);
    }

  hmm trained_model(2, 4);
  trained_model.train(train_data);
  std::cout << "trained model =" << std::endl;
  trained_model.dump();
}

void
hmm_partial_test()
{
  std::cout << "hmm--partial-test ..." << std::endl;
  phmm model(2, 4);
  arma::mat a(2, 2), b(2, 4);
  arma::vec pi(2);

  pi[0] = 0.5;
  pi[1] = 0.5;

  a.at(0, 0) = 0.7;
  a.at(0, 1) = 0.3;
  a.at(1, 0) = 0.2;
  a.at(1, 1) = 0.8;

  b.at(0, 0) = 0.1;
  b.at(0, 1) = 0.2;
  b.at(0, 2) = 0.4;
  b.at(0, 3) = 0.4;
  b.at(1, 0) = 0.5;
  b.at(1, 1) = 0.3;
  b.at(1, 2) = 0.1;
  b.at(1, 3) = 0.1;

  model.setStart(pi);
  model.setTransition(a);
  model.setEmission(b);
  //model.dump();

  phmm::seq_type sseq, p_sseq, oseq, mapseq;
  model.genSeq(3000, sseq, oseq);
  //std::cout << "seq="; hmm::print_seq(sseq, oseq);

  //double prob = model.probability(oseq);
  //std::cout << "logprob=" << prob << " prob=" << exp(prob) << std::endl;

  p_sseq = sseq;
  for (int t = 0; t < (int) p_sseq.size(); t++)
    if (arma::randu() < 0.9) p_sseq[t] = -1;
  model.map(oseq, p_sseq, mapseq);
  //std::cout << "mapseq="; hmm::print_seq(mapseq, sseq);

  double error = 0, total = 0;
  for (unsigned int t = 0; t < sseq.size(); t++)
    {
      if (sseq[t] != mapseq[t] && p_sseq[t] == -1) error++;
      if (p_sseq[t] == -1) total++;
    }
  std::cout << "error=" << error << " = " << error / total << "%" << std::endl;

  hmm model1(2, 4);
  model1.setStart(pi);
  model1.setTransition(a);
  model1.setEmission(b);
  model1.map(oseq, mapseq);
  //std::cout << "mapseq="; hmm::print_seq(mapseq, sseq);
  error = 0;
  for (unsigned int t = 0; t < sseq.size(); t++)
    if (sseq[t] != mapseq[t]) error++;
  std::cout << "error(original)=" << error << " = " << error / sseq.size()
      << "%" << std::endl;

  //  std::vector<phmm::seq_type> train_data;
  //  int N = 1, T = 100;
  //  for (int i = 0; i < N; i++)
  //    {
  //      hmm::seq_type sseq, oseq;
  //      model.genSeq(T, sseq, oseq);
  //      train_data.push_back(oseq);
  //    }
  //
  //  hmm trained_model(2, 4);
  //  trained_model.train(train_data);
  //  std::cout << "trained model =" << std::endl;
  //  trained_model.dump();
}

void
dhmm_test()
{
  std::cout << "dhmm--test ..." << std::endl;
  dhmm model(2, 4);
  arma::mat a(2, 2), b(4, 2);
  arma::vec pi(2);

  pi[0] = 0.5;
  pi[1] = 0.5;

  a.at(0, 0) = 0;
  a.at(1, 0) = 1;

  a.at(0, 1) = 1;
  a.at(1, 1) = 0;

  b.at(0, 0) = 0.1;
  b.at(1, 0) = 0.1;
  b.at(2, 0) = 0.4;
  b.at(3, 0) = 0.4;

  b.at(0, 1) = 0.4;
  b.at(1, 1) = 0.4;
  b.at(2, 1) = 0.1;
  b.at(3, 1) = 0.1;

  model.setStart(pi);
  model.setTransition(a);
  model.setEmission(b);

  dhmm::ivec d;
  dhmm::dvec p;
  d.ones(5);
  d = arma::cumsum(d);
  p.ones(5);
  p /= 5.0;
  model.setDuration(0, d, p);

  d.ones(10);
  d = arma::cumsum(d);
  p.ones(10);
  p /= 10.0;
  model.setDuration(1, d, p);

  //model.dump();

  dhmm::ivec q, o;
  model.genSeq(5, o, q, d); // 5 state-changes
  q.print("q = ");
  d.print("d = ");
  o.print("o = ");

  double prob;

  prob = model.seqProbability(o);
  std::cout << "prob(o) = exp(" << prob << ")" << std::endl;

//  dhmm::ivec qs, ds;
//  model.decode(o, qs, ds);
//  qs.print("qs = ");
//  ds.print("ds = ");
//
  model.train(o);
}

int
main(int argc, char** argv)
{
  process_options(argc, argv);

  if (vm.count("help")) usage();

  //srand(time(NULL));

  //hmm_test();
  //hmm_partial_test();
  dhmm_test();

  return 0;
}

void
process_options(int argc, char** argv)
{
  desc.add_options()("help", "produce this usage")
  //    ("gin", po::value<std::string>(&gin)->default_value(""), "input file (factor graph)")
  //    ("fin", po::value<std::string>(&fin)->default_value(""), "input file (factor functions)")
  //    ("out", po::value<std::string>(&out)->default_value(""), "output file")
  ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}
